import argparse
import pandas as pd
import re

def check_motif(seq, pred_pos_text):
    # 将pred_pos_text转换为整数列表
    pred_pos = list(map(int, pred_pos_text.split(',')))
    
    # 定义motif正则表达式
    motif_patterns = {
        "CG": re.compile(r'CG'),
        "CHH": re.compile(r'C[ACT][ACT]'),
        "CHG": re.compile(r'C[ACT]G')
    }
    
    # 检查每个位置对应的motif
    detected_motifs = set()
    for pos in pred_pos:
        for motif_name, pattern in motif_patterns.items():
            if pos < len(seq) - 2:  # 确保不会超出序列范围
                subseq = seq[pos:pos + 3]
                if pattern.match(subseq):
                    detected_motifs.add(motif_name)
                    break  # 如果一个位置匹配了一个motif，跳过其他motif

    return detected_motifs

def process_tsv(file_path):
    # 读取TSV文件的前100行，没有头文件
    df = pd.read_csv(file_path, sep='\t', header=None, nrows=100)
    
    # 第四列是seq，第六列是pred_pos_text
    seq_col = 3  # pandas的列索引从0开始，因此第四列是索引3
    pred_pos_text_col = 5  # 第六列是索引5
    
    # 检查每行的motif
    for index, row in df.iterrows():
        seq = row[seq_col]
        pred_pos_text = row[pred_pos_text_col]
        motifs_found = check_motif(seq, pred_pos_text)
        print(f"Line {index + 1}: Motifs found: {motifs_found}")

def main():
    # 定义命令行参数
    parser = argparse.ArgumentParser(description="Check for CG, CHH, and CHG motifs in a TSV file.")
    parser.add_argument('file_path', type=str, help='Path to the input TSV file')
    
    # 解析参数
    args = parser.parse_args()
    
    # 处理输入的TSV文件
    process_tsv(args.file_path)

if __name__ == "__main__":
    main()
